"""Example for the Euler equations"""

from hogs.grids import grid1d as grid
from hogs.solvers.euler1d import EELagrangianSolver
from hogs.solvers.flux.euler.lagrangian_flux import LxfFluxLagrangian,\
     GodunovFluxLagrangian, SPHFluxLagrangian

from hogs.solvers.primitive_variable_functions import EEPrimitiveVariable,\
     EEPrimitiveVariableLagrangian

import numpy

# initialize the grid
g = grid.Grid1D()
g.initialize(xlow=-1.0, xhigh=1.0, dx=0.0012, nb=2, nvar=3)

# construct the solver
solver = EELagrangianSolver(gamma=1.4, tf=0.15, nvar=3, grid=g)

# set the flux function
#solver.flux_function = LxfFluxLagrangian()
solver.flux_function = GodunovFluxLagrangian()
#solver.flux_function = flux_lagrangian.SPHFluxLagrangian()

# set the grid for the flux function
solver.flux_function.set_grid( solver.grid )

# primitive variable function
solver.primitive_variable_function = EEPrimitiveVariable(
    gamma=1.4, grid=solver.grid)

# process command line 
solver.setup()

# set the variables
grid = solver.grid

# cell centers
x = grid.xc
q = grid.q

# initial data for the shock tube problem
rhol = 1.0; rhor = 1.0
ul = ur = 0.0

for i, j in enumerate(x):
    if j < 0.6:

        pi = numpy.exp(-50*j*j)

        q[0,i] = rhol
        q[1,i] = ul
        q[2,i] = 2.5 * pi

    else:

        pi = numpy.exp(-50*j*j)

        q[0,i] = rhor
        q[1,i] = ur
        q[2,i] = 2.5 * pi

solver.solve()

# compute the exact solution
# from hogs.solvers.riemann import reuler
# rsolver = reuler.RiemannSolverEulerExact(gamma=1.4)

# pm, um = rsolver.compute_star(rhol, rhor, pl, pr, ul, ur)

# x = solver.grid.xc
# rho = numpy.ones_like(x)
# u = numpy.ones_like(x)
# p = numpy.ones_like(x)

# dx = x[1] - x[0]
# ncells = solver.grid.ncells

# for i in range(ncells):
#     s = (x[i] - 0.6)/0.15

#     _rho, _u, _p = rsolver.sample(pm, um, s, rhol, rhor, pl, pr, ul, ur)

#     rho[i] = _rho; u[i] = _u; p[i] = _p

# numpy.savez('exact', x=x, p=p, rho=rho, u=u)
